{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e0b7c55d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from edm_qm9_utils.analyze import check_stability_yy\n",
    "from edm_qm9_utils.rdkit_functions import build_molecule\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit.Chem.Draw import IPythonConsole\n",
    "from rdkit.Chem import QED\n",
    "from rdkit.Geometry import Point3D\n",
    "import py3Dmol\n",
    "from tqdm import tqdm\n",
    "import pdb\n",
    "from scipy.stats import wasserstein_distance\n",
    "import matplotlib.pyplot as plt\n",
    "from fcd_torch import FCD as FCDMetric\n",
    "from e3fp.pipeline import fprints_from_mol, confs_from_smiles\n",
    "from e3fp.fingerprint.metrics.fprint_metrics import tanimoto\n",
    "import os\n",
    "from frozen_conf_xtb.utils import calculate_XTB_energy\n",
    "import random\n",
    "import utils_yy.utils as utils_yy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8387f13b",
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_atom_list = ['H', 'C', 'O', 'N', 'F']\n",
    "log_dir_ld = './logs/job6_latent_ddpm_qm9/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446ed95e",
   "metadata": {},
   "source": [
    "### 2D valid and unique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f58c2719",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:04<00:00, 2272.98it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:04<00:00, 2247.50it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:04<00:00, 2284.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.         0.95278866 0.95278866] [0.         0.00258223 0.00258223]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "smiless = torch.load(log_dir_ld + 'sample_smiles.pt')\n",
    "\n",
    "res_list = []\n",
    "\n",
    "for nrun in range(3):\n",
    "    num_mol = 0\n",
    "    smiles_sample = []\n",
    "    for smi in tqdm(smiless[nrun*10000:(nrun+1)*10000]):\n",
    "        try:\n",
    "            mol = Chem.MolFromSmiles(smi)\n",
    "            Chem.SanitizeMol(mol)\n",
    "            atom_list = []\n",
    "            for atom in mol.GetAtoms():\n",
    "                atom_list.append(atom.GetSymbol())\n",
    "            if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "                continue\n",
    "\n",
    "            mol = Chem.RemoveHs(mol)\n",
    "            smi = Chem.CanonSmiles( Chem.MolToSmiles(mol) )\n",
    "            smiles_sample.append(smi)\n",
    "        except:\n",
    "            None\n",
    "\n",
    "        num_mol += 1\n",
    "\n",
    "    valid = len(smiles_sample) / num_mol\n",
    "    unique = len(set(smiles_sample)) / len(smiles_sample)\n",
    "    valid_and_unique = valid * unique\n",
    "#     print(valid, unique, valid_and_unique, num_mol)\n",
    "    res_list.append([valid, unique, valid_and_unique])\n",
    "\n",
    "res_list = np.array(res_list)\n",
    "print(res_list.mean(axis=0), res_list.std(axis=0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bdfc96e",
   "metadata": {},
   "source": [
    "### 3D stable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f8de29d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [01:18<00:00, 127.42it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [01:21<00:00, 122.41it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [01:22<00:00, 121.52it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.86873996 0.97522809] [0.00230675 0.00025238]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "res_list = []\n",
    "\n",
    "for nrun in range(3):\n",
    "    valid_num = 0\n",
    "    total_num = 0\n",
    "    atom_valid = 0\n",
    "    for mol in tqdm(mols[nrun*10000:(nrun+1)*10000]):    \n",
    "        atom_list = []\n",
    "        for atom in mol.GetAtoms():\n",
    "            atom_list.append(atom.GetSymbol())\n",
    "        if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "            continue\n",
    "\n",
    "        AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "        positions = mol.GetConformers()[0].GetPositions()\n",
    "\n",
    "        out = check_stability_yy(positions, atom_list)\n",
    "\n",
    "        if out[0]:\n",
    "            valid_num += 1\n",
    "        total_num += 1\n",
    "\n",
    "        atom_valid += out[1] / out[2]\n",
    "\n",
    "#     print(valid_num / total_num, atom_valid / total_num, valid_num, total_num)\n",
    "    res_list.append([valid_num / total_num, atom_valid / total_num])\n",
    "\n",
    "res_list = np.array(res_list)\n",
    "print(res_list.mean(axis=0), res_list.std(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d14b57d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:16<00:00, 623.36it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:16<00:00, 620.27it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:16<00:00, 610.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00229521 0.3261557 ] [0.00017812 0.00240093]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "res_list = []\n",
    "\n",
    "for nrun in range(3):\n",
    "    valid_num = 0\n",
    "    total_num = 0\n",
    "    atom_valid = 0\n",
    "    for mol in tqdm(mols[nrun*10000:(nrun+1)*10000]):    \n",
    "        atom_list = []\n",
    "        for atom in mol.GetAtoms():\n",
    "            atom_list.append(atom.GetSymbol())\n",
    "        if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "            continue\n",
    "\n",
    "#         AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "        positions = mol.GetConformers()[0].GetPositions()\n",
    "\n",
    "        out = check_stability_yy(positions, atom_list)\n",
    "\n",
    "        if out[0]:\n",
    "            valid_num += 1\n",
    "        total_num += 1\n",
    "\n",
    "        atom_valid += out[1] / out[2]\n",
    "\n",
    "#     print(valid_num / total_num, atom_valid / total_num, valid_num, total_num)\n",
    "    res_list.append([valid_num / total_num, atom_valid / total_num])\n",
    "\n",
    "res_list = np.array(res_list)\n",
    "print(res_list.mean(axis=0), res_list.std(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ac23672e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [08:40<00:00, 19.22it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [07:41<00:00, 21.65it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [08:57<00:00, 18.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.82011158 0.89047233] [0.00367722 0.00348876]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "res_list = []\n",
    "\n",
    "for nrun in range(3):\n",
    "    valid_num = 0\n",
    "    total_num = 0\n",
    "    atom_valid = 0\n",
    "    for mol in tqdm(mols[nrun*10000:(nrun+1)*10000]):    \n",
    "        atom_list = []\n",
    "        for atom in mol.GetAtoms():\n",
    "            atom_list.append(atom.GetSymbol())\n",
    "        if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "            continue\n",
    "\n",
    "        try:\n",
    "            AllChem.EmbedMolecule(mol)\n",
    "            AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "            positions = mol.GetConformers()[0].GetPositions()\n",
    "        except:\n",
    "            total_num += 1\n",
    "            continue\n",
    "\n",
    "        out = check_stability_yy(positions, atom_list)\n",
    "\n",
    "        if out[0]:\n",
    "            valid_num += 1\n",
    "        total_num += 1\n",
    "\n",
    "        atom_valid += out[1] / out[2]\n",
    "\n",
    "#     print(valid_num / total_num, atom_valid / total_num, valid_num, total_num)\n",
    "    res_list.append([valid_num / total_num, atom_valid / total_num])\n",
    "\n",
    "res_list = np.array(res_list)\n",
    "print(res_list.mean(axis=0), res_list.std(axis=0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7e0780a",
   "metadata": {},
   "source": [
    "### 2D distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "78001ee2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 100000/100000 [00:00<00:00, 295250.22it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[    0     2     5     7    25    91   475  2404 13625 83366]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "data = np.load('/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules/qm9/latent_diffusion/emb_2d_3d/train.npz')\n",
    "\n",
    "stats = np.zeros(10, dtype=int)\n",
    "for ncharge in tqdm(data['charges']):\n",
    "    idx = len([n for n in ncharge if n > 1])\n",
    "    stats[idx] += 1\n",
    "print(stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f57a4669",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████| 100000/100000 [00:00<00:00, 4933372.54it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████| 13083/13083 [00:00<00:00, 5293659.97it/s]\n"
     ]
    }
   ],
   "source": [
    "with open('/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules/data/smiles_qm9.txt', 'r') as f:\n",
    "    smiles_train = [smi for smi in tqdm(f.read().split('\\n')[:100000])]\n",
    "    \n",
    "with open('/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules/data/smiles_qm9.txt', 'r') as f:\n",
    "    smiles_test = [smi for smi in tqdm(f.read().split('\\n')[-13084:-1])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a5a37d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "def measure_2d_distribution(smi_list1, smi_list2, _3fold=True):\n",
    "    qed_list1 = []\n",
    "    for smi in tqdm(smi_list1):\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        prop = QED.properties(mol)\n",
    "        qed = QED.qed(mol)\n",
    "        qed_list1.append([prop[0], prop[1], prop[4], qed])\n",
    "    random.shuffle(qed_list1)\n",
    "    qed_list1 = np.array(qed_list1)\n",
    "    \n",
    "    qed_list2 = []\n",
    "    for smi in tqdm(smi_list2):\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        prop = QED.properties(mol)\n",
    "        qed = QED.qed(mol)\n",
    "        qed_list2.append([prop[0], prop[1], prop[4], qed])\n",
    "    qed_list2 = np.array(qed_list2)\n",
    "    \n",
    "    # print results\n",
    "    prop_name = ['MW', 'ALogP', 'PSA', 'QED']\n",
    "    if _3fold:\n",
    "        for nprop, pname in enumerate(prop_name):\n",
    "            num_run = qed_list1.shape[0] // 3\n",
    "            res_list = [[v for v in utils_yy.statistical_metric(qed_list1[:num_run, nprop],\n",
    "                                                                qed_list2[:, nprop]).values()],\n",
    "                        [v for v in utils_yy.statistical_metric(qed_list1[num_run:num_run*2, nprop],\n",
    "                                                                qed_list2[:, nprop]).values()],\n",
    "                        [v for v in utils_yy.statistical_metric(qed_list1[num_run*2:, nprop],\n",
    "                                                                qed_list2[:, nprop]).values()]]\n",
    "\n",
    "            key_list = [k for k in utils_yy.statistical_metric(qed_list1[:num_run, nprop],\n",
    "                                                               qed_list2[:, nprop]).keys()]\n",
    "\n",
    "            print(pname)\n",
    "            print({k: [np.mean([res_list[0][n], res_list[1][n], res_list[2][n]]),\n",
    "                       np.std([res_list[0][n], res_list[1][n], res_list[2][n]])]\n",
    "                   for n, k in enumerate(key_list)})\n",
    "        \n",
    "        kwargs_fcd = {'n_jobs': 8, 'device': 'cpu', 'batch_size': 32}\n",
    "        fcd = [ FCDMetric(**kwargs_fcd)(smi_list1[:num_run], smi_list2),\n",
    "                FCDMetric(**kwargs_fcd)(smi_list1[num_run:num_run*2], smi_list2),\n",
    "                FCDMetric(**kwargs_fcd)(smi_list1[num_run*2:], smi_list2) ]\n",
    "        print('FCD', np.mean(fcd), np.std(fcd))\n",
    "        \n",
    "    else:\n",
    "        for nprop, pname in enumerate(prop_name):\n",
    "            print(pname)\n",
    "            print(utils_yy.statistical_metric(qed_list1[:, nprop], qed_list2[:, nprop]))\n",
    "        \n",
    "        kwargs_fcd = {'n_jobs': 8, 'device': 'cpu', 'batch_size': 32}\n",
    "        print('FCD', FCDMetric(**kwargs_fcd)(smi_list1, smi_list2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "67a03e0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████| 100000/100000 [02:20<00:00, 713.06it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 13083/13083 [00:18<00:00, 700.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MW\n",
      "{'tvd': 0.006152899946495427, 'hd': 0.027250938652787673, 'wd': 0.11524130402583446}\n",
      "ALogP\n",
      "{'tvd': 0.002624401131239011, 'hd': 0.030700064941371456, 'wd': 0.007492805439065966}\n",
      "PSA\n",
      "{'tvd': 0.0036061522586562714, 'hd': 0.029461469455488783, 'wd': 0.14610340198731134}\n",
      "QED\n",
      "{'tvd': 0.0030609500879003257, 'hd': 0.03105649267647254, 'wd': 0.0005733954490561486}\n",
      "FCD 0.0298101927620813\n"
     ]
    }
   ],
   "source": [
    "measure_2d_distribution(smiles_train, smiles_test, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f4a60335",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 27501/27501 [00:38<00:00, 705.72it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 13083/13083 [00:17<00:00, 732.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MW\n",
      "{'tvd': [0.02935428198405649, 0.004752470157714284], 'hd': [0.17141575439215453, 0.0008324221259389271], 'wd': [1.7714501981004018, 0.044561127593093615]}\n",
      "ALogP\n",
      "{'tvd': [0.008299265308101352, 0.0005354909008047986], 'hd': [0.08299838748179272, 0.0006707623465969748], 'wd': [0.17530174062687146, 0.005110986789145081]}\n",
      "PSA\n",
      "{'tvd': [0.023748301545992025, 0.00201212106210784], 'hd': [0.11327890115482468, 0.0024406406175092672], 'wd': [2.611887956084219, 0.09649659085273403]}\n",
      "QED\n",
      "{'tvd': [0.007445369333508467, 0.0002900820798815914], 'hd': [0.06541844392446715, 0.0021318058924504326], 'wd': [0.005765394189978828, 0.0005135423771765456]}\n",
      "FCD 0.580426556853895 0.003947453378448896\n"
     ]
    }
   ],
   "source": [
    "log_dir_edm = '/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules_official/qm9/'\n",
    "with open(log_dir_edm + 'generate_smiles.txt', 'r') as f:\n",
    "    smiless = f.read().split('\\n')[:-1]\n",
    "\n",
    "measure_2d_distribution(smiless, smiles_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e5e39ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████| 100000/100000 [00:25<00:00, 3995.41it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 30004/30004 [00:40<00:00, 733.17it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 13083/13083 [00:17<00:00, 727.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MW\n",
      "{'tvd': [0.022450151808060984, 0.0012644474612296625], 'hd': [0.08116829252875644, 0.00167361944333405], 'wd': [0.5999122989745161, 0.027839864635935553]}\n",
      "ALogP\n",
      "{'tvd': [0.009378349441438382, 0.0015421847509588763], 'hd': [0.06997382708192325, 0.0011263225847387278], 'wd': [0.10555170970283016, 0.001756792043358978]}\n",
      "PSA\n",
      "{'tvd': [0.012172530110945755, 0.0001771692903032922], 'hd': [0.09651223941333957, 0.0030863856017845164], 'wd': [1.925151165889883, 0.044115261258926834]}\n",
      "QED\n",
      "{'tvd': [0.010089493592951167, 6.627141043087967e-05], 'hd': [0.11516015716413623, 0.0020413791495492743], 'wd': [0.018175715857238813, 0.0006957501238402865]}\n"
     ]
    }
   ],
   "source": [
    "smiless = torch.load(log_dir_ld + 'sample_smiles.pt')\n",
    "smiless1 = []\n",
    "\n",
    "stats_tar = stats / stats.sum() * 30000\n",
    "stats_gen = np.zeros(stats.shape)\n",
    "\n",
    "for smi in tqdm(smiless):\n",
    "    try:\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        mol = Chem.RemoveHs(mol)\n",
    "        atom_list = []\n",
    "        for atom in mol.GetAtoms():\n",
    "            atom_list.append(atom.GetSymbol())\n",
    "        if len(atom_list) > 9:\n",
    "            continue\n",
    "        if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "            continue\n",
    "        \n",
    "        num_atom = len(atom_list)\n",
    "        if stats_gen[num_atom] >= stats_tar[num_atom]:\n",
    "            continue\n",
    "        \n",
    "        stats_gen[num_atom] += 1\n",
    "        smiless1.append(smi)\n",
    "    except:\n",
    "        None\n",
    "\n",
    "measure_2d_distribution(smiless1, smiles_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "49eab202",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▋                                                                                                | 689/100000 [00:00<00:42, 2334.51it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ERROR! Session/line number was not unique in database. History logging moved to new session 119\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|█████████████████████████████████████████▌                                                     | 43713/100000 [00:11<00:14, 3824.82it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 30001/30001 [00:43<00:00, 690.00it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 13083/13083 [00:19<00:00, 672.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MW\n",
      "{'tvd': [0.06503261255239606, 0.0025482457597286695], 'hd': [0.2591391819198361, 0.00309431895076475], 'wd': [5.4630820434476, 0.12172579652712512]}\n",
      "ALogP\n",
      "{'tvd': [0.008884374195768529, 0.0011423928813224448], 'hd': [0.06434184958920419, 0.0034629449615702183], 'wd': [0.07050763718947335, 0.006751313857472962]}\n",
      "PSA\n",
      "{'tvd': [0.01054676109881432, 0.00024042704314579822], 'hd': [0.09394433453907099, 0.00395845558034042], 'wd': [1.3967515074453392, 0.11232934920778169]}\n",
      "QED\n",
      "{'tvd': [0.011007745975892871, 0.0011655775891214397], 'hd': [0.07328189839646003, 0.0014192370764311085], 'wd': [0.007266553271805258, 0.0005465462436851693]}\n",
      "FCD 0.4428336272238518 0.009556058317030347\n"
     ]
    }
   ],
   "source": [
    "smiless = torch.load(log_dir_ld + 'sample_smiles.pt')\n",
    "smiless1 = []\n",
    "\n",
    "stats_tar = stats / stats.sum() * 30000\n",
    "stats_gen = np.zeros(stats.shape)\n",
    "\n",
    "for smi in tqdm(smiless):\n",
    "    try:\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        mol = Chem.RemoveHs(mol)\n",
    "        atom_list = []\n",
    "        for atom in mol.GetAtoms():\n",
    "            atom_list.append(atom.GetSymbol())\n",
    "        if len(atom_list) > 9:\n",
    "            continue\n",
    "        if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "            continue\n",
    "        \n",
    "#         num_atom = len(atom_list)\n",
    "#         if stats_gen[num_atom] >= stats_tar[num_atom]:\n",
    "#             continue\n",
    "        \n",
    "        stats_gen[num_atom] += 1\n",
    "        smiless1.append(smi)\n",
    "        \n",
    "        if len(smiless1) > 30000:\n",
    "            break\n",
    "    except:\n",
    "        None\n",
    "\n",
    "measure_2d_distribution(smiless1, smiles_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9db6092",
   "metadata": {},
   "source": [
    "### 3D distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "533c9597",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['num_atoms', 'charges', 'positions', 'index', 'A', 'B', 'C', 'mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'U0', 'U', 'H', 'G', 'Cv', 'omega1', 'emb_2d', 'emb_3d', 'zpve_thermo', 'U0_thermo', 'U_thermo', 'H_thermo', 'G_thermo', 'Cv_thermo']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 13083/13083 [06:49<00:00, 31.98it/s]\n"
     ]
    }
   ],
   "source": [
    "data = np.load('/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules/qm9/latent_diffusion/emb_2d_3d/test.npz')\n",
    "print([k for k in data.keys()])\n",
    "\n",
    "dataset_info = {'name': 'qm9', 'atom_decoder': {1:'H', 6:'C', 8:'O', 7:'N', 9:'F'}}\n",
    "energy_list = []\n",
    "for idx, natoms in enumerate(tqdm(data['num_atoms'])):\n",
    "    x = data['charges'][idx][:natoms]\n",
    "    pos = data['positions'][idx][:natoms]\n",
    "    \n",
    "    geometry = []\n",
    "    for i, p in zip(x, pos):\n",
    "        geometry.append([dataset_info['atom_decoder'][i]] + p.tolist())\n",
    "    \n",
    "    energy = calculate_XTB_energy(geometry)\n",
    "    energy_list.append(energy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2fb5dea4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tvd': 0.003885490330963845, 'hd': 0.029642486119174343, 'wd': 0.031156386598832922}\n"
     ]
    }
   ],
   "source": [
    "# data = np.load('/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules/qm9/latent_diffusion/emb_2d_3d/train.npz')\n",
    "# print([k for k in data.keys()])\n",
    "\n",
    "# dataset_info = {'name': 'qm9', 'atom_decoder': {1:'H', 6:'C', 8:'O', 7:'N', 9:'F'}}\n",
    "# energy_train_list = []\n",
    "# for idx, natoms in enumerate(tqdm(data['num_atoms'])):\n",
    "#     x = data['charges'][idx][:natoms]\n",
    "#     pos = data['positions'][idx][:natoms]\n",
    "    \n",
    "#     geometry = []\n",
    "#     for i, p in zip(x, pos):\n",
    "#         geometry.append([dataset_info['atom_decoder'][i]] + p.tolist())\n",
    "    \n",
    "#     energy = calculate_XTB_energy(geometry)\n",
    "#     energy_train_list.append(energy)\n",
    "\n",
    "# torch.save(energy_train_list, './qm9_training_data_energy.pt')\n",
    "\n",
    "\n",
    "energy_train_list = torch.load('./qm9_training_data_energy.pt')\n",
    "\n",
    "print(utils_yy.statistical_metric(np.array(energy_train_list), np.array(energy_list)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "97b81013",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30000/30000 [28:52<00:00, 17.31it/s]\n"
     ]
    }
   ],
   "source": [
    "mol_dir = '/scratch/user/yuning.you/project/graph_latent_diffusion/e3_diffusion_for_molecules_official/outputs/edm_qm9/eval/analyzed_molecules/'\n",
    "atom2idx = {'H':1, 'C':6, 'O':8, 'N':7, 'F':9}\n",
    "energy_edm_list = []\n",
    "for fn in tqdm(os.listdir(mol_dir)):\n",
    "    x, pos = [], []\n",
    "    try:\n",
    "        with open(mol_dir + fn, 'r') as f:\n",
    "            data = f.read().split('\\n')\n",
    "    except:\n",
    "        continue\n",
    "    \n",
    "    natoms = int(data[0])\n",
    "    for d in data[2:2+natoms]:\n",
    "        d = d.split()\n",
    "        x.append(d[0])\n",
    "        pos.append([float(d[1]), float(d[2]), float(d[3])])\n",
    "    \n",
    "    geometry = []\n",
    "    for i, p in zip(x, pos):\n",
    "        geometry.append([i] + p)\n",
    "    \n",
    "    energy = calculate_XTB_energy(geometry)\n",
    "    energy_edm_list.append(energy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "199e638f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tvd': [0.02810979388009885, 0.002992548617162737], 'hd': [0.13542332908722657, 0.0035131040141895818], 'wd': [0.28568282655895605, 0.0045918625204167465]}\n"
     ]
    }
   ],
   "source": [
    "random.shuffle(energy_edm_list)\n",
    "num_run = len(energy_edm_list) // 3\n",
    "res_list = [[v for v in utils_yy.statistical_metric(np.array(energy_edm_list[:num_run]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_edm_list[num_run:num_run*2]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_edm_list[num_run*2:]),\n",
    "                                                    np.array(energy_list)).values()]]\n",
    "\n",
    "key_list = [k for k in utils_yy.statistical_metric(np.array(energy_edm_list[:num_run]),\n",
    "                                                   np.array(energy_list)).keys()]\n",
    "\n",
    "print({k: [np.mean([res_list[0][n], res_list[1][n], res_list[2][n]]),\n",
    "           np.std([res_list[0][n], res_list[1][n], res_list[2][n]])]\n",
    "       for n, k in enumerate(key_list)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "96052f17",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100000/100000 [16:48<00:00, 99.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tvd': [0.018782709851518756, 0.0018861222341144625], 'hd': [0.08252747344333826, 0.0024023886386782692], 'wd': [0.25082174310225563, 0.0026850363893648897]}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "stats_tar = stats / stats.sum() * 30000\n",
    "stats_gen = np.zeros(stats.shape)\n",
    "\n",
    "energy_ld_list = []\n",
    "for mol in tqdm(mols):    \n",
    "    atom_list = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        atom_list.append(atom.GetSymbol())\n",
    "#     if len(atom_list) > 29:\n",
    "#         continue\n",
    "    if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "        continue\n",
    "    \n",
    "    AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "    positions = mol.GetConformers()[0].GetPositions()\n",
    "\n",
    "    geometry = []\n",
    "    for idx, a in enumerate(atom_list):\n",
    "        geometry.append([a] + positions[idx].tolist())\n",
    "\n",
    "    mol = Chem.RemoveHs(mol)\n",
    "    num_atom = len(mol.GetAtoms())\n",
    "    if num_atom > 9 or stats_gen[num_atom] >= stats_tar[num_atom]:\n",
    "        continue\n",
    "    stats_gen[num_atom] += 1\n",
    "\n",
    "    energy = calculate_XTB_energy(geometry)\n",
    "    energy_ld_list.append(energy)\n",
    "\n",
    "\n",
    "random.shuffle(energy_ld_list)\n",
    "num_run = len(energy_ld_list) // 3\n",
    "res_list = [[v for v in utils_yy.statistical_metric(np.array(energy_ld_list[:num_run]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_ld_list[num_run:num_run*2]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_ld_list[num_run*2:]),\n",
    "                                                    np.array(energy_list)).values()]]\n",
    "\n",
    "key_list = [k for k in utils_yy.statistical_metric(np.array(energy_ld_list[:num_run]),\n",
    "                                                   np.array(energy_list)).keys()]\n",
    "\n",
    "print({k: [np.mean([res_list[0][n], res_list[1][n], res_list[2][n]]),\n",
    "           np.std([res_list[0][n], res_list[1][n], res_list[2][n]])]\n",
    "       for n, k in enumerate(key_list)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e3a748d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100000/100000 [07:14<00:00, 230.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tvd': [0.7739614360899444, 0.04337661685806737], 'hd': [0.7251906559774723, 0.03224364822916765], 'wd': [1015.7221504497578, 10.111491648653107]}\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "stats_tar = stats / stats.sum() * 30000\n",
    "stats_gen = np.zeros(stats.shape)\n",
    "\n",
    "energy_ld_list = []\n",
    "for mol in tqdm(mols):    \n",
    "    atom_list = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        atom_list.append(atom.GetSymbol())\n",
    "#     if len(atom_list) > 29:\n",
    "#         continue\n",
    "    if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "        continue\n",
    "    \n",
    "#     AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "    positions = mol.GetConformers()[0].GetPositions()\n",
    "\n",
    "    geometry = []\n",
    "    for idx, a in enumerate(atom_list):\n",
    "        geometry.append([a] + positions[idx].tolist())\n",
    "\n",
    "    mol = Chem.RemoveHs(mol)\n",
    "    num_atom = len(mol.GetAtoms())\n",
    "    if num_atom > 9 or stats_gen[num_atom] >= stats_tar[num_atom]:\n",
    "        continue\n",
    "    stats_gen[num_atom] += 1\n",
    "\n",
    "    try:\n",
    "        energy = calculate_XTB_energy(geometry)\n",
    "    except:\n",
    "        energy = None\n",
    "    energy_ld_list.append(energy)\n",
    "\n",
    "energy_mean = np.mean([e for e in energy_ld_list if e is not None])\n",
    "energy_ld_list = [e if e is not None else energy_mean for e in energy_ld_list]\n",
    "\n",
    "random.shuffle(energy_ld_list)\n",
    "num_run = len(energy_ld_list) // 3\n",
    "res_list = [[v for v in utils_yy.statistical_metric(np.array(energy_ld_list[:num_run]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_ld_list[num_run:num_run*2]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_ld_list[num_run*2:]),\n",
    "                                                    np.array(energy_list)).values()]]\n",
    "\n",
    "key_list = [k for k in utils_yy.statistical_metric(np.array(energy_ld_list[:num_run]),\n",
    "                                                   np.array(energy_list)).keys()]\n",
    "\n",
    "print({k: [np.mean([res_list[0][n], res_list[1][n], res_list[2][n]]),\n",
    "           np.std([res_list[0][n], res_list[1][n], res_list[2][n]])]\n",
    "       for n, k in enumerate(key_list)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d6f9e3ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 73%|██████████████████████████████████████████████████████████████████████▊                          | 72980/100000 [15:12<05:37, 79.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tvd': [0.04053214469936226, 0.0061896470352884635], 'hd': [0.0979752837054297, 0.0018371192676422707], 'wd': [0.28276628173420465, 0.014545023241275275]}\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "stats_tar = stats / stats.sum() * 30000\n",
    "stats_gen = np.zeros(stats.shape)\n",
    "\n",
    "energy_rand_list = []\n",
    "for mol in tqdm(mols):    \n",
    "    atom_list = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        atom_list.append(atom.GetSymbol())\n",
    "#     if len(atom_list) > 29:\n",
    "#         continue\n",
    "    if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "        continue\n",
    "    \n",
    "    try:\n",
    "        mol = Chem.RemoveHs(mol)\n",
    "        num_atom = len(mol.GetAtoms())\n",
    "        if num_atom > 9 or stats_gen[num_atom] >= stats_tar[num_atom]:\n",
    "            continue\n",
    "        stats_gen[num_atom] += 1\n",
    "        mol = Chem.AddHs(mol)\n",
    "        \n",
    "        AllChem.EmbedMolecule(mol)\n",
    "        AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "        positions = mol.GetConformers()[0].GetPositions()\n",
    "\n",
    "        geometry = []\n",
    "        for idx, a in enumerate(atom_list):\n",
    "            geometry.append([a] + positions[idx].tolist())\n",
    "\n",
    "        energy = calculate_XTB_energy(geometry)\n",
    "    \n",
    "    except:\n",
    "        energy = None\n",
    "\n",
    "    energy_rand_list.append(energy)\n",
    "    if len(energy_rand_list) > 30000:\n",
    "        break\n",
    "\n",
    "energy_mean = np.mean([e for e in energy_rand_list if e is not None])\n",
    "energy_rand_list = [e if e is not None else energy_mean for e in energy_rand_list]\n",
    "\n",
    "random.shuffle(energy_rand_list)\n",
    "num_run = len(energy_rand_list) // 3\n",
    "res_list = [[v for v in utils_yy.statistical_metric(np.array(energy_rand_list[:num_run]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_rand_list[num_run:num_run*2]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_rand_list[num_run*2:]),\n",
    "                                                    np.array(energy_list)).values()]]\n",
    "\n",
    "key_list = [k for k in utils_yy.statistical_metric(np.array(energy_rand_list[:num_run]),\n",
    "                                                   np.array(energy_list)).keys()]\n",
    "\n",
    "print({k: [np.mean([res_list[0][n], res_list[1][n], res_list[2][n]]),\n",
    "           np.std([res_list[0][n], res_list[1][n], res_list[2][n]])]\n",
    "       for n, k in enumerate(key_list)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b9186661",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|██████████████████████████████████████████▍                                                      | 43713/100000 [09:55<12:46, 73.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tvd': [0.03497220961029582, 0.003946201454430411], 'hd': [0.19081330710655867, 0.003650139484239148], 'wd': [1.060405469194736, 0.01896135130480013]}\n"
     ]
    }
   ],
   "source": [
    "mols = torch.load(log_dir_ld + 'sample_conformer.pt')\n",
    "\n",
    "stats_tar = stats / stats.sum() * 30000\n",
    "stats_gen = np.zeros(stats.shape)\n",
    "\n",
    "energy_ld_list = []\n",
    "for mol in tqdm(mols):    \n",
    "    atom_list = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        atom_list.append(atom.GetSymbol())\n",
    "#     if len(atom_list) > 29:\n",
    "#         continue\n",
    "    if len(set(atom_list).difference(qm9_atom_list)) > 0:\n",
    "        continue\n",
    "    \n",
    "    AllChem.MMFFOptimizeMolecule(mol, maxIters=200)\n",
    "    positions = mol.GetConformers()[0].GetPositions()\n",
    "\n",
    "    geometry = []\n",
    "    for idx, a in enumerate(atom_list):\n",
    "        geometry.append([a] + positions[idx].tolist())\n",
    "\n",
    "    mol = Chem.RemoveHs(mol)\n",
    "    num_atom = len(mol.GetAtoms())\n",
    "    if num_atom > 9:\n",
    "        continue\n",
    "#     if num_atom > 9 or stats_gen[num_atom] >= stats_tar[num_atom]:\n",
    "#         continue\n",
    "#     stats_gen[num_atom] += 1\n",
    "\n",
    "    energy = calculate_XTB_energy(geometry)\n",
    "    energy_ld_list.append(energy)\n",
    "    \n",
    "    if len(energy_ld_list) > 30000:\n",
    "        break\n",
    "\n",
    "\n",
    "random.shuffle(energy_ld_list)\n",
    "num_run = len(energy_ld_list) // 3\n",
    "res_list = [[v for v in utils_yy.statistical_metric(np.array(energy_ld_list[:num_run]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_ld_list[num_run:num_run*2]),\n",
    "                                                    np.array(energy_list)).values()],\n",
    "            [v for v in utils_yy.statistical_metric(np.array(energy_ld_list[num_run*2:]),\n",
    "                                                    np.array(energy_list)).values()]]\n",
    "\n",
    "key_list = [k for k in utils_yy.statistical_metric(np.array(energy_ld_list[:num_run]),\n",
    "                                                   np.array(energy_list)).keys()]\n",
    "\n",
    "print({k: [np.mean([res_list[0][n], res_list[1][n], res_list[2][n]]),\n",
    "           np.std([res_list[0][n], res_list[1][n], res_list[2][n]])]\n",
    "       for n, k in enumerate(key_list)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7db736d5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
